{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DaSiamRPN.net import SiamRPNBIG\n",
    "from DaSiamRPN.utils import get_axis_aligned_bbox, cxy_wh_2_rect\n",
    "from DaSiamRPN.run_SiamRPN import SiamRPN_init, SiamRPN_track\n",
    "import cv2 as cv\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "import random\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from gen_seq import gen_seq\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DaSiamRPN():\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def load_model(self, net_file):\n",
    "        self.net = SiamRPNBIG()\n",
    "        self.net.load_state_dict(torch.load(net_file, map_location=\"cpu\"))\n",
    "        self.net.eval().cpu()\n",
    "\n",
    "        for i in range(10):\n",
    "            self.net.temple(torch.autograd.Variable(\n",
    "                torch.FloatTensor(1, 3, 127, 127)).cpu())\n",
    "            self.net(torch.autograd.Variable(\n",
    "                torch.FloatTensor(1, 3, 255, 255)).cpu())\n",
    "\n",
    "    def init_bbox(self, image, bbox):\n",
    "        target_pos = np.array(bbox[0:2])\n",
    "        target_size = np.array(bbox[2:4])\n",
    "        self.state = SiamRPN_init(image, target_pos, target_size, self.net)\n",
    "\n",
    "    def update(self, image):\n",
    "        self.state = SiamRPN_track(self.state, image)\n",
    "        target_pos = self.state['target_pos']\n",
    "        target_size = self.state['target_sz']\n",
    "        bbox = [target_pos[0], target_pos[1], target_size[0], target_size[1]]\n",
    "        return bbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_rect(image, rect, color):\n",
    "    x = int(rect[0])\n",
    "    y = int(rect[1])\n",
    "    width = int(rect[2])\n",
    "    height = int(rect[3])\n",
    "\n",
    "    if color == \"green\":\n",
    "        color = (0, 255, 0)\n",
    "    elif color == \"red\":\n",
    "        color = (255, 0, 0)\n",
    "    elif color == \"blue\":\n",
    "        color = (0, 0, 255)\n",
    "\n",
    "    cv.rectangle(image, (x, y), (x+width, y+height), color)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Generate sequence config\n",
    "    img_list, init_bbox, gt = gen_seq(seq='Football')\n",
    "\n",
    "    # Tracker\n",
    "    tracker = DaSiamRPN()\n",
    "    tracker.load_model(\"./DaSiamRPN/SiamRPNBIG.model\")\n",
    "\n",
    "    # init\n",
    "    image = cv.imread(img_list[0])\n",
    "    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)\n",
    "    \n",
    "    tracker.init_bbox(image, gt[0])\n",
    "    \n",
    "    dpi = 80.0\n",
    "    figsize = (image.shape[0] / dpi, image.shape[1] / dpi)\n",
    "    fig = plt.figure(frameon=False, figsize=figsize, dpi=dpi)\n",
    "    ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
    "    ax.set_axis_off()\n",
    "    fig.add_axes(ax)\n",
    "    im = ax.imshow(image)\n",
    "\n",
    "    # Run tracker\n",
    "    for img_file, ground_truth in zip(img_list[1:4], gt[1:4]):\n",
    "        image = cv.imread(img_file)\n",
    "        image = cv.cvtColor(image, cv.COLOR_BGR2RGB)\n",
    "\n",
    "        # draw init\n",
    "        image_draw = image.copy()\n",
    "        image_draw = draw_rect(image_draw, ground_truth, \"red\")\n",
    "\n",
    "        # update\n",
    "        result_bbox = tracker.update(image)\n",
    "        print(result_bbox)\n",
    "\n",
    "        image_draw = draw_rect(image_draw, result_bbox, \"blue\")\n",
    "\n",
    "        # image_draw = cv.cvtColor(image_draw, cv.COLOR_RGB2BGR)\n",
    "        # cv.imshow(\"image\", image_draw)\n",
    "        # cv.waitKey(0)\n",
    "        # plt.imshow(image_draw)\n",
    "        ax.imshow(image_draw)\n",
    "        plt.show()\n",
    "\n",
    "    # Save result\n",
    "    # res = {}\n",
    "    # res['res'] = result_bb.round().tolist()\n",
    "    # res['type'] = 'rect'\n",
    "    # res['fps'] = fps\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
